use crate::block::Block;
use crate::consts::{BIN_BLOCK_SIZE, PTR_SIZE};
use crate::utils::access_guard::AccessGuard;
use crate::utils::atomic_single_linked_list::AtomicSingleLinkedList;
use crate::utils::double_linked_list::DoubleLinkedNode;
use crate::utils::single_linked_list::{SingleLinkedList, SingleLinkedNode};
use std::alloc::Layout;
use std::mem;
use std::ptr::NonNull;
use std::sync::atomic::{AtomicBool, Ordering};

#[derive(Eq, PartialEq)]
pub(crate) enum PageKind {
    Small, // small blocks go into 64kb pages inside a segment

    // Unsupported for now
    #[allow(dead_code)]
    Large, // larger blocks go into a single page spanning a whole segment
    #[allow(dead_code)]
    Huge, // huge blocks (>512kb) are put into a single page in a segment of the exact size (but still 2mb aligned)
}
#[derive(Debug)]
pub(crate) struct Page {
    bin: usize,
    length: usize,
    used: usize,
    in_sleep: bool,
    in_delayed: AtomicBool,
    page_area: NonNull<u8>,

    free: SingleLinkedList<Block>,
    local_free: SingleLinkedList<Block>,
    thread_free: AtomicSingleLinkedList<Block>,

    next: Option<NonNull<Page>>,
    prev: Option<NonNull<Page>>,
    sleep_next: Option<NonNull<Page>>,
}
impl Page {
    pub(crate) fn init(
        page_ptr: NonNull<Page>,
        bin: usize,
        page_size: usize,
        page_area: NonNull<u8>,
    ) {
        debug_assert!(PTR_SIZE * BIN_BLOCK_SIZE[bin] <= page_size);
        debug_assert!(
            page_ptr.as_ptr() as usize + size_of::<Self>() <= page_area.as_ptr() as usize
        );
        let page = Self {
            bin,
            length: page_size / (PTR_SIZE * BIN_BLOCK_SIZE[bin]),
            used: 0,
            in_sleep: false,
            in_delayed: AtomicBool::new(false),
            page_area,
            free: Default::default(),
            local_free: Default::default(),
            thread_free: Default::default(),
            next: None,
            prev: None,
            sleep_next: None,
        };
        unsafe {
            page_ptr.write(page);
            page_ptr.as_ref().init_all_blocks();
        }
    }
}

impl Page {
    fn init_all_blocks(&self) {
        (0..self.length)
            .map(|idx| {
                // 计算 block 的指针地址
                // 指针一定不为空
                unsafe {
                    self.page_area
                        .byte_add(idx * PTR_SIZE * BIN_BLOCK_SIZE[self.bin])
                }
            })
            .for_each(|block_ptr| {
                let block_ptr = block_ptr.cast();
                self.access(block_ptr).init();
                self.free.push(self, block_ptr)
            });
    }
    pub(crate) fn bin(&self) -> usize {
        self.bin
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.used == 0
    }

    pub(crate) fn is_full(&self) -> bool {
        self.used == self.length
    }

    pub(crate) fn set_in_sleep(&mut self, in_sleep: bool) {
        self.in_sleep = in_sleep;
    }

    pub(crate) fn is_in_sleep(&self) -> bool {
        self.in_sleep
    }

    /// 这个函数涉及跨线程操作，需要保证原子性
    ///
    /// 返回值为是否是当前线程修改成功，多线程并发竞争时，应该有且只有一个线程能返回 true
    pub(crate) fn set_in_delayed(&mut self, in_delayed: bool) -> bool {
        self.in_delayed
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |x| {
                if x == in_delayed {
                    None
                } else {
                    Some(in_delayed)
                }
            })
            .is_ok()
    }

    pub(crate) fn is_in_delayed(&self) -> bool {
        self.in_delayed.load(Ordering::Relaxed)
    }

    pub(crate) fn alloc(&mut self, layout: Layout) -> Option<NonNull<Block>> {
        debug_assert!(
            PTR_SIZE * BIN_BLOCK_SIZE[self.bin] == 0
                || PTR_SIZE * BIN_BLOCK_SIZE[self.bin] >= layout.size()
        );
        self.free
            .pop(self)
            .or_else(|| {
                self.local_collect();
                self.free.pop(self)
            })
            .or_else(|| {
                self.thread_collect();
                self.free.pop(self)
            })
            .and_then(|block| {
                self.used += 1;
                Some(block)
            })
    }

    pub(crate) fn dealloc(&mut self, block: NonNull<Block>, local: bool) {
        let block = self.access(block);
        block.init();
        if local {
            // 同线程 free 块，不涉及跨线程问题
            self.local_free.push(self, block);
            self.used -= 1;
        } else {
            // 跨线程 free 块
            self.thread_free.push(self, block)
        }
    }

    /// 收集由本地线程释放的块，接入 free 列表中
    /// 这个操作非常轻，可以直接在快过程中使用
    fn local_collect(&mut self) {
        debug_assert!(self.free.is_empty());
        mem::swap(&mut self.local_free, &mut self.free);
    }

    /// 收集由其他地线程释放的块，接入 free 列表中
    /// 这个操作需要遍历 thread_free 链表
    pub(crate) fn thread_collect(&mut self) {
        debug_assert!(self.free.is_empty());
        // local_free 无可用块，尝试收集 thread_free 可用快
        let first_free_block = self.thread_free.take();
        self.used -= first_free_block.iter(self).count();
        self.free = first_free_block;
    }
}

impl AccessGuard<Block> for Page {
    fn before_access(&self, ptr: NonNull<Block>) {
        // 检查 block 指针处于 page 范围内
        // 检查 block 指针是否符合对齐
        let ptr = ptr.as_ptr() as usize;
        let page_area_ptr = self.page_area.as_ptr() as usize;
        debug_assert!(ptr >= page_area_ptr);
        debug_assert_eq!(
            (ptr - page_area_ptr) % (PTR_SIZE * BIN_BLOCK_SIZE[self.bin]),
            0
        );
        let idx = (ptr - page_area_ptr) / (PTR_SIZE * BIN_BLOCK_SIZE[self.bin]);
        debug_assert!(idx < self.length);
    }
}

impl DoubleLinkedNode<Self> for Page {
    fn next(&mut self) -> &mut Option<NonNull<Self>> {
        &mut self.next
    }
    fn prev(&mut self) -> &mut Option<NonNull<Self>> {
        &mut self.prev
    }
}

impl SingleLinkedNode<Self> for Page {
    fn next(&mut self) -> &mut Option<NonNull<Self>> {
        &mut self.sleep_next
    }
}

#[cfg(test)]
mod tests {
    use crate::consts::{BIN_BLOCK_SIZE, PTR_SIZE, SMALL_PAGE_SHIFT};
    use crate::heap::Heap;
    use crate::page::{Page, PageKind};
    use crate::segment::Segment;
    use crate::utils::access_guard::AccessGuard;
    use crate::utils::single_linked_list::SingleLinkedNode;
    use std::alloc::Layout;
    use std::ptr::NonNull;

    fn assert_page(page_ptr: NonNull<Page>) {
        Heap::thread_local_with(|heap| {
            let page = heap.access(page_ptr);
            // 检查块一定存在于三个链表之中一个，并且块总数一致
            // 块指针合法，块指针不重复
            let free = page.free.iter(page).collect::<Vec<_>>();
            let local_free = page.local_free.iter(page).collect::<Vec<_>>();
            let mut thread_free = vec![];
            let mut current = page.thread_free.first();
            while let Some(block) = current {
                thread_free.push(block);
                let block = page.access(block);
                current = *block.next();
            }
            let mut all_blocks = free
                .iter()
                .chain(local_free.iter())
                .chain(thread_free.iter())
                .collect::<Vec<_>>();
            all_blocks.dedup_by(|a, b| a.as_ptr() == b.as_ptr());
            assert_eq!(all_blocks.len() + page.used, page.length);
        });
    }
    #[test]
    fn test_page_basic() {
        const TEST_BIN: usize = 24;
        const TEST_ALLOC_BYTES: usize = BIN_BLOCK_SIZE[TEST_BIN] * PTR_SIZE;
        let (seg_owner, free_pages) =
            Segment::new(0, PageKind::Small, TEST_BIN, SMALL_PAGE_SHIFT).expect("Out of memory");
        let page_ptr = free_pages.first().unwrap();
        let seg = &seg_owner.0;
        free_pages.give_up_all(seg.as_ref());
        let page = seg.access(page_ptr);
        println!("Get page in test, page: {page:?}");
        assert_page(page_ptr);
        assert_eq!(page.bin(), TEST_BIN);
        assert!(page.is_empty());

        // 分配并释放一个块
        let layout = Layout::from_size_align(TEST_ALLOC_BYTES, TEST_ALLOC_BYTES).unwrap();
        let block = page.alloc(layout).unwrap();
        assert_page(page_ptr);
        page.dealloc(block, true);
        assert_page(page_ptr);

        // 全部分配
        let mut blocks = vec![];
        while let Some(b) = page.alloc(layout) {
            blocks.push(b);
            assert_page(page_ptr);
        }
        assert_page(page_ptr);
        assert!(page.alloc(layout).is_none());
        assert_eq!(blocks.len(), page.length);

        // 全部归还
        blocks.iter().for_each(|b| {
            page.dealloc(*b, true);
            assert_page(page_ptr);
        });

        assert!(page.free.is_empty());
        assert!(page.is_empty());

        // 再次申请，local_free 链表应该直接转移到 free 上
        let block = page.alloc(layout).unwrap();
        assert_page(page_ptr);
        assert!(!page.free.is_empty());
        assert!(page.local_free.is_empty());
        page.dealloc(block, true);

        // 必须确保 seg_owner 在最后 drop，否则前面都是悬垂指针。
        drop(seg_owner);
    }
}
